﻿
/*Created by jiadong chen
 * https://jiadong-chen.medium.com/
 *用来烘焙动作贴图。烘焙对象使用Animation组件，并且在导入时设置Rig为Legacy
 */
#if UNITY_EDITOR
using System.Collections.Generic;
using System.Linq;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.UIElements;
using Color = UnityEngine.Color;

/// <summary>
/// 保存需要烘焙的动画的相关数据
/// </summary>
public struct AnimData
{
    #region FIELDS

    private int _vertexCount;
    private int _mapWidth;
    private readonly List<AnimationState> _animClips;
    private string _name;

    private Animation _animation;
    private AnimationClip setupPose;
    private float setupPoseTime;
    private SkinnedMeshRenderer _skin;
    public Dictionary<Transform, int3> _meshindex;
    public Transform _transform;

    public List<AnimationState> AnimationClips => _animClips;
    public int MapWidth => _mapWidth;
    public string Name => _name;

    #endregion

    public AnimData(Animation anim, SkinnedMeshRenderer smr, GameObject go, Dictionary<Transform, int3> meshindex, AnimationClip spc, float spt)
    {
        _vertexCount = smr.sharedMesh.vertexCount;
        _mapWidth = Mathf.NextPowerOfTwo(_vertexCount);
        _animClips = new List<AnimationState>(anim.Cast<AnimationState>());
        _animation = anim;
        _skin = smr;
        _name = go.name;
        _meshindex = meshindex;
        _transform = go.transform;
        setupPose = spc;
        setupPoseTime = spt;
    }

    #region METHODS

    public void AnimationPlay(string animName)
    {
        SetupPose();
        _animation.Play(animName);
        _animation.GetComponent<Animator>().Play(animName);
    }
    public void SetupPose()
    {
        if (setupPose != null)
        {
            setupPose.SampleAnimation(_transform.gameObject, setupPoseTime);
        }
    }
    public void SampleAnimAndBakeMesh(float dt, ref Mesh m)
    {
        SampleAnim(dt);
        BakeMesh(ref m);
    }

    private void SampleAnim(float dt)
    {
        if (_animation == null)
        {
            Debug.LogError("animation is null!!");
            return;
        }
        _animation.GetComponent<Animator>().Update(dt);
        _animation.Sample();
    }

    private void BakeMesh(ref Mesh m)
    {
        if (_skin == null)
        {
            Debug.LogError("skin is null!!");
            return;
        }

        _skin.BakeMesh(m);
    }


    #endregion

}

/// <summary>
/// 烘焙后的数据
/// </summary>
public struct BakedData
{
    #region FIELDS

    private readonly string _name;
    private readonly float _animLen;
    private readonly byte[] _rawAnimMap;
    private readonly int _animMapWidth;
    private readonly int _animMapHeight;

    #endregion

    public BakedData(string name, float animLen, Texture2D animMap)
    {
        _name = name;
        _animLen = animLen;
        _animMapHeight = animMap.height;
        _animMapWidth = animMap.width;
        _rawAnimMap = animMap.GetRawTextureData();
    }

    public int AnimMapWidth => _animMapWidth;

    public string Name => _name;

    public float AnimLen => _animLen;

    public byte[] RawAnimMap => _rawAnimMap;

    public int AnimMapHeight => _animMapHeight;
}

/// <summary>
/// 烘焙器
/// </summary>
public class AnimMapBaker
{

    #region FIELDS

    private AnimData? _animData = null;
    private Mesh _bakedMesh;
    private readonly List<Vector3> _vertices = new List<Vector3>();
    private readonly List<BakedData> _bakedDataList = new List<BakedData>();

    #endregion

    #region METHODS

    public void SetAnimData(GameObject go, Dictionary<Transform, int3> meshindex, AnimationClip sp, float spt)
    {
        if (go == null)
        {
            Debug.LogError("go is null!!");
            return;
        }

        var anim = go.GetComponent<Animation>();
        var smr = go.GetComponentInChildren<SkinnedMeshRenderer>();

        if (anim == null || smr == null)
        {
            Debug.LogError("anim or smr is null!!");
            return;
        }
        _bakedMesh = new Mesh();
        _animData = new AnimData(anim, smr, go, meshindex, sp, spt);
    }

    /// <summary>
    /// merge为true表示输出结果为一张图
    /// </summary>
    /// <param name="merge"></param>
    /// <returns></returns>
    public List<BakedData> Bake(bool merge)
    {
        if (_animData == null)
        {
            Debug.LogError("bake data is null!!");
            return _bakedDataList;
        }

        //每一个动作都生成一个动作图
        if (!merge)
        {
            foreach (var t in _animData.Value.AnimationClips)
            {
                if (!t.clip.legacy)
                {
                    Debug.LogError(string.Format($"{t.clip.name} is not legacy!!"));
                    continue;
                }
                BakePerAnimClip(t);
            }
        }
        //合并为一张动作图
        else
        {
            List<AnimationState> _clips = new List<AnimationState>();
            foreach (var t in _animData.Value.AnimationClips)
            {
                if (!t.clip.legacy)
                {
                    Debug.LogError(string.Format($"{t.clip.name} is not legacy!!"));
                    continue;
                }
                else
                    _clips.Add(t);
            }
            BakePerAnimClip(_clips);
        }

        return _bakedDataList;
    }
    private void BakePerAnimClip(List<AnimationState> curAnims)
    {
        var i = 0;
        var curClipFrame = 0;
        Texture2D animMap;
        int totalFrame = 0;
        float totalTime = 0;
        foreach (var item in curAnims)
        {
            totalFrame += Mathf.ClosestPowerOfTwo((int)(item.clip.frameRate * item.length));
            totalTime += item.length;
        }
        animMap = new Texture2D(_animData.Value.MapWidth * 2, totalFrame + 1, TextureFormat.RGBAHalf, true);
        animMap.name = string.Format($"{_animData.Value.Name}.animMap");
        for (int c = 0; c < curAnims.Count; c++)
        {
            AnimationState curAnim = curAnims[c];


            float sampleTime = 0;
            float perFrameTime = 0;

            int cur = Mathf.ClosestPowerOfTwo((int)(curAnim.clip.frameRate * curAnim.length));
            Color animdata = new Color();
            animdata.g = (float)curClipFrame / totalFrame;
            curClipFrame += cur;
            perFrameTime = curAnim.length / cur;
            _animData.Value.AnimationPlay(curAnim.name);

            Color[] colorEvent = new Color[_animData.Value.MapWidth];
            float[] blend = BlendTypes(_animData.Value.MapWidth);

            for (; i < curClipFrame; i++)
            {
                curAnim.time = sampleTime;

                _animData.Value.SampleAnimAndBakeMesh(perFrameTime, ref _bakedMesh);
                ApplyColor(ref colorEvent);
                for (var j = 0; j < _bakedMesh.vertexCount * 2; j += 2)
                {
                    var vertex = _bakedMesh.vertices[j / 2];
                    animMap.SetPixel(j, i + 1, new Color(vertex.x, vertex.y, vertex.z, blend[j / 2]));

                    animMap.SetPixel(j + 1, i + 1, colorEvent[j / 2]);
                }

                sampleTime += perFrameTime;
            }
            animdata.r = curAnim.length;
            animdata.b = (float)curClipFrame / totalFrame;
            animMap.SetPixel(c, 0, animdata);//将动画在AnimMap上的位置信息储存在第一行中
            curAnim.clip.legacy = false;
        }
        animMap.Apply(true, true);
        _animData.Value.SetupPose();
        _bakedDataList.Add(new BakedData(animMap.name, totalTime, animMap));
    }
    private void BakePerAnimClip(AnimationState curAnim)
    {
        var curClipFrame = 0;
        float sampleTime = 0;
        float perFrameTime = 0;

        curClipFrame = Mathf.ClosestPowerOfTwo((int)(curAnim.clip.frameRate * curAnim.length));
        perFrameTime = curAnim.length / curClipFrame;

        var animMap = new Texture2D(_animData.Value.MapWidth * 2, curClipFrame, TextureFormat.RGBAHalf, true);
        animMap.name = string.Format($"{_animData.Value.Name}_{curAnim.name}.animMap");
        _animData.Value.AnimationPlay(curAnim.name);

        Color[] colorEvent = new Color[_animData.Value.MapWidth];
        float[] blend = BlendTypes(_animData.Value.MapWidth);

        for (var i = 0; i < curClipFrame; i++)
        {
            curAnim.time = sampleTime;


            _animData.Value.SampleAnimAndBakeMesh(perFrameTime, ref _bakedMesh);
            ApplyColor(ref colorEvent);
            for (var j = 0; j < _bakedMesh.vertexCount * 2; j += 2)
            {
                var vertex = _bakedMesh.vertices[j / 2];
                animMap.SetPixel(j, i + 1, new Color(vertex.x, vertex.y, vertex.z, blend[j / 2]));

                animMap.SetPixel(j + 1, i + 1, colorEvent[j / 2]);
            }

            sampleTime += perFrameTime;
        }
        animMap.Apply();
        _animData.Value.SetupPose();
        _bakedDataList.Add(new BakedData(animMap.name, curAnim.clip.length, animMap));
    }

    private float[] BlendTypes(int size)
    {
        float[] blentT = new float[size];
        foreach (var item in _animData.Value._meshindex)
        {
            for (int i = item.Value.x; i <= item.Value.y; i++)
            {
                blentT[i] = item.Value.z;
            }
        }
        return blentT;
    }
    private void ApplyColor(ref Color[] colorEvent)
    {
        MaterialPropertyBlock mpb = new MaterialPropertyBlock();
        foreach (var item in _animData.Value._meshindex)
        {
            var skmr = item.Key.GetComponent<SkinnedMeshRenderer>();
            Color color;
            skmr.GetPropertyBlock(mpb);
            color = item.Key.gameObject.activeSelf ? (skmr.HasPropertyBlock() ? mpb.GetColor("_Color") : new Color(1, 1, 1, 1)) : new Color(0, 0, 0, -1);
            for (int i = item.Value.x; i <= item.Value.y; i++)
            {
                colorEvent[i] = color;
            }
        }
    }

    #endregion

}
#endif
